"""
author : DengXiuqi
date : 2018.10
email : dengxiuqi@163.com
"""

from config import *
import tensorflow as tf
from network import SRGAN_g
from data import ImgData
from utils import saveImg
import os


def train(reload=True):
    # set up database
    train_db = ImgData(batch_size, 'train')
    test_db = ImgData(batch_size, 'test')

    image = tf.placeholder(dtype=tf.float32, shape=[None, img_size[0], img_size[1], 3], name='image')
    target = tf.placeholder(dtype=tf.float32, shape=[None, img_size[0], img_size[1], 3], name='target')

    '''Loss'''
    net_g = SRGAN_g(image, is_train=True, reuse=False)
    mse_loss = tf.reduce_sum(tf.square(net_g - target), name='gm')
    loss = mse_loss

    '''Optimizer'''
    optimizer = tf.train.AdamOptimizer(learn_rate).minimize(loss)
    saver = tf.train.Saver(max_to_keep=10)
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_proportion, allow_growth=True)
    config = tf.ConfigProto(gpu_options=gpu_options)

    '''Train'''
    with tf.Session(config=config) as sess:
        sess.run(tf.global_variables_initializer())
        step, loss_total = start_step, 0

        print("is training")
        model_dir = "/home/tree/dxq/model"
        if not os.path.exists(model_dir):
            # make /model direction
            os.mkdir(model_dir)
            print("create the directory: %s" % model_dir)

        if reload:
            checkPoint = tf.train.get_checkpoint_state(model_dir)
            # if have checkPoint, restore checkPoint
            if checkPoint and checkPoint.model_checkpoint_path:
                saver.restore(sess, checkPoint.model_checkpoint_path)
                print("restored %s" % checkPoint.model_checkpoint_path)
            else:
                print("no checkpoint found!")

        while train_db.epoch < total_epoch:
            # train
            img_clear, img_blur = train_db.next_batch()
            _, _l = sess.run([optimizer, loss], feed_dict={image: img_blur, target: img_clear})
            loss_total += _l

            if step % 100 == 0 and step != start_step:
                # print the loss every 100 step
                print("epoch: ", train_db.epoch, " step: ", step, "loss: ", loss_total / (1 if not step else 100))
                loss_total = 0

            if step % save_step == 0 and step != start_step:
                # print the loss every save_step(int) step
                img_clear, img_blur = test_db.next_batch()
                res = sess.run(net_g, feed_dict={image: img_blur, target: img_clear})
                saveImg(res, test_db.cursor-batch_size, name='step%d_' % step)
                saver.save(sess, './model/model-', global_step=step)
            step += 1


if __name__ == "__main__":
    train()


